/**
 * @file session.h
 * @brief 定义了 txnn.cpp 推理会话管理类.
 */

#pragma once

#include <memory>
#include <string>

#include "utils/data/image_data.h"
#include "infer_config.h"
#include "manager/manager.h"
#include "manager/tensor.h"

namespace infer {

/**
 * @class TxnnSession 执行推理流程类
 * @brief 封装后端推理的子流程, 子流程封装了后端推理的启动、执行、结束等操作
 */
class TxnnSession {
public:
  /**
   * @brief TxnnSession构造函数,初始化TxnnSession对象.
   *
   * @param engine_path Txnn编译后的文件目录.
   * @param cfg 推理配置文件.
   */
  TxnnSession(const std::string& engine_path, const InferConfig& cfg);

  TxnnSession() = default;
  TxnnSession(const std::string& name);
  ~TxnnSession();

  /**
   * @brief 执行推理操作。
   *
   * @param inputs 输入图像向量。
   */
  void infer(const std::vector<image::Image>& inputs);

  // txdaStream_t stream_;      ///< 推理流: TODO 暂时不可用
  InferConfig cfg_;             ///< 推理配置
  std::vector<Tensor> tensors;  ///< 张量列表
  std::vector<dim_t> shape_min; ///< 输入张量最小形状: B,H,W,C (TBD)
  std::vector<dim_t> shape_max; ///< 输入张量最大形状: B,H,W,C (TBD)
  bool is_dynamic_batch_;       ///< 是否为动态batch

private:
  
  void init();
  void getTensorInfo();
  void inferSingleBatch(const std::vector<image::Image>& inputs);
  void inferDynamicBatch(const std::vector<image::Image>& inputs);
  
  std::unique_ptr<TxnnManager> manager_;

  // TODO 整图推理 或 多图推理

  bool zero_copy_; ///< 是否使用零拷贝

  int input_shape_; ///< 输入张量形状
  int infer_shape_; ///< 推理张量形状
};

} // namespace infer